/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.apache.flink.table.utils;

import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.typeutils.TupleTypeInfo;
import org.apache.flink.table.api.DataTypes;
import org.apache.flink.table.api.ValidationException;
import org.apache.flink.table.catalog.ResolvedSchema;
import org.apache.flink.table.legacy.api.TableSchema;
import org.apache.flink.table.legacy.sinks.TableSink;
import org.apache.flink.table.legacy.sources.DefinedProctimeAttribute;
import org.apache.flink.table.legacy.sources.DefinedRowtimeAttributes;
import org.apache.flink.table.legacy.sources.RowtimeAttributeDescriptor;
import org.apache.flink.table.legacy.sources.TableSource;
import org.apache.flink.table.types.DataType;
import org.apache.flink.table.types.logical.LegacyTypeInformationType;
import org.apache.flink.table.types.utils.DataTypeUtils;
import org.apache.flink.table.types.utils.TypeConversions;
import org.apache.flink.types.Row;

import org.junit.jupiter.api.Test;

import javax.annotation.Nullable;

import java.util.Collections;
import java.util.List;
import java.util.function.Function;
import java.util.stream.Collectors;

import static org.apache.flink.table.api.DataTypes.BIGINT;
import static org.apache.flink.table.api.DataTypes.DECIMAL;
import static org.apache.flink.table.api.DataTypes.FIELD;
import static org.apache.flink.table.api.DataTypes.ROW;
import static org.apache.flink.table.api.DataTypes.STRING;
import static org.apache.flink.table.api.DataTypes.TIME;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;

/** Tests for {@link TypeMappingUtils}. */
class TypeMappingUtilsTest {

    @Test
    void testFieldMappingReordered() {
        int[] indices =
                TypeMappingUtils.computePhysicalIndices(
                        TableSchema.builder()
                                .field("f1", DataTypes.BIGINT())
                                .field("f0", DataTypes.STRING())
                                .build()
                                .getTableColumns(),
                        ROW(FIELD("f0", DataTypes.STRING()), FIELD("f1", DataTypes.BIGINT())),
                        Function.identity());

        assertThat(indices).isEqualTo(new int[] {1, 0});
    }

    @Test
    void testFieldMappingNonMatchingTypes() {
        assertThatThrownBy(
                        () ->
                                TypeMappingUtils.computePhysicalIndices(
                                        TableSchema.builder()
                                                .field("f1", DataTypes.BIGINT())
                                                .field("f0", DataTypes.TIMESTAMP(3))
                                                .build()
                                                .getTableColumns(),
                                        ROW(
                                                FIELD("f0", DataTypes.STRING()),
                                                FIELD("f1", DataTypes.BIGINT())),
                                        Function.identity()))
                .isInstanceOf(ValidationException.class)
                .hasMessage(
                        "Type TIMESTAMP(3) of table field 'f0' does not match with the physical "
                                + "type STRING of the 'f0' field of the TableSource return type.");
    }

    @Test
    void testFieldMappingNonMatchingPrecision() {
        assertThatThrownBy(
                        () ->
                                TypeMappingUtils.computePhysicalIndices(
                                        TableSchema.builder()
                                                .field("f0", DataTypes.TIMESTAMP(9))
                                                .build()
                                                .getTableColumns(),
                                        ROW(FIELD("f0", DataTypes.TIMESTAMP(3))),
                                        Function.identity()))
                .isInstanceOf(ValidationException.class)
                .hasMessage(
                        "Type TIMESTAMP(9) of table field 'f0' does not match with the physical "
                                + "type TIMESTAMP(3) of the 'f0' field of the TableSource "
                                + "return type.");
    }

    @Test
    void testNameMappingDoesNotExist() {
        assertThatThrownBy(
                        () ->
                                TypeMappingUtils.computePhysicalIndices(
                                        TableSchema.builder()
                                                .field("f0", DataTypes.BIGINT())
                                                .build()
                                                .getTableColumns(),
                                        ROW(FIELD("f0", DataTypes.BIGINT())),
                                        str -> null))
                .isInstanceOf(ValidationException.class)
                .hasMessage("Field 'f0' could not be resolved by the field mapping.");
    }

    @Test
    void testFieldMappingLegacyDecimalType() {
        int[] indices =
                TypeMappingUtils.computePhysicalIndices(
                        TableSchema.builder()
                                .field("f0", DECIMAL(38, 18))
                                .build()
                                .getTableColumns(),
                        ROW(FIELD("f0", TypeConversions.fromLegacyInfoToDataType(Types.BIG_DEC))),
                        Function.identity());

        assertThat(indices).isEqualTo(new int[] {0});
    }

    @Test
    void testFieldMappingLegacyDecimalTypeNotMatchingPrecision() {
        assertThatThrownBy(
                        () ->
                                TypeMappingUtils.computePhysicalIndices(
                                        TableSchema.builder()
                                                .field("f0", DECIMAL(38, 10))
                                                .build()
                                                .getTableColumns(),
                                        ROW(
                                                FIELD(
                                                        "f0",
                                                        TypeConversions.fromLegacyInfoToDataType(
                                                                Types.BIG_DEC))),
                                        Function.identity()))
                .isInstanceOf(ValidationException.class)
                .hasMessage(
                        "Type DECIMAL(38, 10) of table field 'f0' does not match with the "
                                + "physical type LEGACY('DECIMAL', 'DECIMAL') of the "
                                + "'f0' field of the TableSource return type.")
                .cause()
                .isInstanceOf(ValidationException.class)
                .hasMessage("Legacy decimal type can only be mapped to DECIMAL(38, 18).");
    }

    @Test
    void testFieldMappingRowTypeNotMatchingNamesInNestedType() {
        int[] indices =
                TypeMappingUtils.computePhysicalIndices(
                        TableSchema.builder()
                                .field("f0", DECIMAL(38, 18))
                                .field(
                                        "f1",
                                        ROW(
                                                FIELD("logical_f1_0", BIGINT()),
                                                FIELD("logical_f1_1", STRING())))
                                .build()
                                .getTableColumns(),
                        ROW(
                                FIELD("f0", DECIMAL(38, 18)),
                                FIELD(
                                        "f1",
                                        ROW(
                                                FIELD("physical_f1_0", BIGINT()),
                                                FIELD("physical_f1_1", STRING())))),
                        Function.identity());

        assertThat(indices).isEqualTo(new int[] {0, 1});
    }

    @Test
    void testFieldMappingRowTypeNotMatchingTypesInNestedType() {
        assertThatThrownBy(
                        () ->
                                TypeMappingUtils.computePhysicalIndices(
                                        TableSchema.builder()
                                                .field("f0", DECIMAL(38, 18))
                                                .field(
                                                        "f1",
                                                        ROW(
                                                                FIELD("f1_0", BIGINT()),
                                                                FIELD("f1_1", STRING())))
                                                .build()
                                                .getTableColumns(),
                                        ROW(
                                                FIELD("f0", DECIMAL(38, 18)),
                                                FIELD(
                                                        "f1",
                                                        ROW(
                                                                FIELD("f1_0", STRING()),
                                                                FIELD("f1_1", STRING())))),
                                        Function.identity()))
                .isInstanceOf(ValidationException.class)
                .hasMessage(
                        "Type ROW<`f1_0` BIGINT, `f1_1` STRING> of table field 'f1' does not "
                                + "match with the physical type ROW<`f1_0` STRING, `f1_1` STRING> "
                                + "of the 'f1' field of the TableSource return type.");
    }

    @Test
    void testFieldMappingLegacyCompositeType() {
        int[] indices =
                TypeMappingUtils.computePhysicalIndices(
                        TableSchema.builder()
                                .field("f1", DataTypes.BIGINT())
                                .field("f0", DataTypes.STRING())
                                .build()
                                .getTableColumns(),
                        TypeConversions.fromLegacyInfoToDataType(
                                Types.TUPLE(Types.STRING, Types.LONG)),
                        Function.identity());

        assertThat(indices).isEqualTo(new int[] {1, 0});
    }

    @Test
    void testFieldMappingLegacyCompositeTypeWithRenaming() {
        int[] indices =
                TypeMappingUtils.computePhysicalIndices(
                        TableSchema.builder()
                                .field("a", DataTypes.BIGINT())
                                .field("b", DataTypes.STRING())
                                .build()
                                .getTableColumns(),
                        TypeConversions.fromLegacyInfoToDataType(
                                Types.TUPLE(Types.STRING, Types.LONG)),
                        str -> {
                            switch (str) {
                                case "a":
                                    return "f1";
                                case "b":
                                    return "f0";
                                default:
                                    throw new AssertionError();
                            }
                        });

        assertThat(indices).isEqualTo(new int[] {1, 0});
    }

    @Test
    void testMappingWithBatchTimeAttributes() {
        TestTableSource tableSource =
                new TestTableSource(
                        DataTypes.BIGINT(), Collections.singletonList("rowtime"), "proctime");
        int[] indices =
                TypeMappingUtils.computePhysicalIndicesOrTimeAttributeMarkers(
                        tableSource,
                        TableSchema.builder()
                                .field("a", Types.LONG)
                                .field("rowtime", Types.SQL_TIMESTAMP)
                                .field("proctime", Types.SQL_TIMESTAMP)
                                .build()
                                .getTableColumns(),
                        false,
                        Function.identity());

        assertThat(indices).isEqualTo(new int[] {0, -3, -4});
    }

    @Test
    void testMappingWithStreamTimeAttributes() {
        TestTableSource tableSource =
                new TestTableSource(
                        DataTypes.BIGINT(), Collections.singletonList("rowtime"), "proctime");
        int[] indices =
                TypeMappingUtils.computePhysicalIndicesOrTimeAttributeMarkers(
                        tableSource,
                        TableSchema.builder()
                                .field("a", Types.LONG)
                                .field("rowtime", Types.SQL_TIMESTAMP)
                                .field("proctime", Types.SQL_TIMESTAMP)
                                .build()
                                .getTableColumns(),
                        true,
                        Function.identity());

        assertThat(indices).isEqualTo(new int[] {0, -1, -2});
    }

    @Test
    void testMappingWithStreamTimeAttributesFromCompositeType() {
        TestTableSource tableSource =
                new TestTableSource(
                        ROW(FIELD("b", TIME()), FIELD("a", DataTypes.BIGINT())),
                        Collections.singletonList("rowtime"),
                        "proctime");
        int[] indices =
                TypeMappingUtils.computePhysicalIndicesOrTimeAttributeMarkers(
                        tableSource,
                        TableSchema.builder()
                                .field("a", Types.LONG)
                                .field("rowtime", Types.SQL_TIMESTAMP)
                                .field("proctime", Types.SQL_TIMESTAMP)
                                .build()
                                .getTableColumns(),
                        true,
                        Function.identity());

        assertThat(indices).isEqualTo(new int[] {1, -1, -2});
    }

    @Test
    void testWrongLogicalTypeForRowtimeAttribute() {
        TestTableSource tableSource =
                new TestTableSource(
                        DataTypes.BIGINT(), Collections.singletonList("rowtime"), "proctime");

        assertThatThrownBy(
                        () ->
                                TypeMappingUtils.computePhysicalIndicesOrTimeAttributeMarkers(
                                        tableSource,
                                        TableSchema.builder()
                                                .field("a", Types.LONG)
                                                .field("rowtime", Types.SQL_TIME)
                                                .field("proctime", Types.SQL_TIMESTAMP)
                                                .build()
                                                .getTableColumns(),
                                        false,
                                        Function.identity()))
                .isInstanceOf(ValidationException.class)
                .hasMessage(
                        "Rowtime field 'rowtime' has invalid type TIME(0). Rowtime attributes "
                                + "must be of a Timestamp family.");
    }

    @Test
    void testWrongLogicalTypeForProctimeAttribute() {
        TestTableSource tableSource =
                new TestTableSource(
                        DataTypes.BIGINT(), Collections.singletonList("rowtime"), "proctime");
        assertThatThrownBy(
                        () ->
                                TypeMappingUtils.computePhysicalIndicesOrTimeAttributeMarkers(
                                        tableSource,
                                        TableSchema.builder()
                                                .field("a", Types.LONG)
                                                .field("rowtime", Types.SQL_TIMESTAMP)
                                                .field("proctime", Types.SQL_TIME)
                                                .build()
                                                .getTableColumns(),
                                        false,
                                        Function.identity()))
                .isInstanceOf(ValidationException.class)
                .hasMessage(
                        "Proctime field 'proctime' has invalid type TIME(0). Proctime attributes "
                                + "must be of a Timestamp family.");
    }

    @Test
    void testCheckPhysicalLogicalTypeCompatible() {
        TableSchema tableSchema =
                TableSchema.builder()
                        .field("a", DataTypes.VARCHAR(2))
                        .field("b", DataTypes.DECIMAL(20, 2))
                        .build();
        TableSink tableSink = new TestTableSink(tableSchema);
        LegacyTypeInformationType legacyDataType =
                (LegacyTypeInformationType) tableSink.getConsumedDataType().getLogicalType();
        TypeInformation legacyTypeInfo =
                ((TupleTypeInfo) legacyDataType.getTypeInformation()).getTypeAt(1);
        DataType physicalType = TypeConversions.fromLegacyInfoToDataType(legacyTypeInfo);
        ResolvedSchema physicSchema = DataTypeUtils.expandCompositeTypeToSchema(physicalType);
        DataType[] logicalDataTypes = tableSchema.getFieldDataTypes();
        List<DataType> physicalDataTypes = physicSchema.getColumnDataTypes();
        for (int i = 0; i < logicalDataTypes.length; i++) {
            TypeMappingUtils.checkPhysicalLogicalTypeCompatible(
                    physicalDataTypes.get(i).getLogicalType(),
                    logicalDataTypes[i].getLogicalType(),
                    "physicalField",
                    "logicalField",
                    false);
        }
    }

    private static class TestTableSource
            implements TableSource<Object>, DefinedProctimeAttribute, DefinedRowtimeAttributes {

        private final DataType producedDataType;
        private final List<String> rowtimeAttributes;
        private final String proctimeAttribute;

        private TestTableSource(
                DataType producedDataType,
                List<String> rowtimeAttributes,
                String proctimeAttribute) {
            this.producedDataType = producedDataType;
            this.rowtimeAttributes = rowtimeAttributes;
            this.proctimeAttribute = proctimeAttribute;
        }

        @Nullable
        @Override
        public String getProctimeAttribute() {
            return proctimeAttribute;
        }

        @Override
        public List<RowtimeAttributeDescriptor> getRowtimeAttributeDescriptors() {
            return rowtimeAttributes.stream()
                    .map(attr -> new RowtimeAttributeDescriptor(attr, null, null))
                    .collect(Collectors.toList());
        }

        @Override
        public DataType getProducedDataType() {
            return producedDataType;
        }

        @Override
        public TableSchema getTableSchema() {
            throw new UnsupportedOperationException("Should not be called");
        }
    }

    /**
     * Since UpsertStreamTableSink not in flink-table-common module, here we use Tuple2 &lt;Boolean,
     * Row&gt; to simulate the behavior of UpsertStreamTableSink.
     */
    private static class TestTableSink implements TableSink<Tuple2<Boolean, Row>> {
        private final TableSchema tableSchema;

        private TestTableSink(TableSchema tableSchema) {
            this.tableSchema = tableSchema;
        }

        TypeInformation<Row> getRecordType() {
            return tableSchema.toRowType();
        }

        @Override
        public TypeInformation<Tuple2<Boolean, Row>> getOutputType() {
            return new TupleTypeInfo<>(Types.BOOLEAN, getRecordType());
        }

        @Override
        public String[] getFieldNames() {
            return tableSchema.getFieldNames();
        }

        @Override
        public TypeInformation<?>[] getFieldTypes() {
            return tableSchema.getFieldTypes();
        }

        @Override
        public TableSink<Tuple2<Boolean, Row>> configure(
                String[] fieldNames, TypeInformation<?>[] fieldTypes) {
            return null;
        }
    }
}
